
use std::rc::Rc;
use std::cell::RefCell;

#[derive(Debug, PartialEq, Eq)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

struct Solution;

impl Solution {
//    // 自己解法
//    pub fn prune_tree(root: Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
//        // 如果全树没有0节点
//        if root.is_some() && !Solution::traverse_tree(root.clone()) {
//            return None;
//        }
//        root
//    }
//
//    // 是否包含1
//    fn traverse_tree(root: Option<Rc<RefCell<TreeNode>>>) -> bool {
//        if let Some(root_rc) = root {
//            let left_flag = Solution::traverse_tree(root_rc.borrow().left.clone());
//            if !left_flag {
//                root_rc.borrow_mut().left = None;
//            }
//            let right_flag = Solution::traverse_tree(root_rc.borrow().right.clone());
//            if !right_flag {
//                root_rc.borrow_mut().right = None;
//            }
//            let mid_flag = root_rc.borrow().val == 1;
//            return left_flag || right_flag || mid_flag;
//        }
//        false
//    }

    pub fn prune_tree(root: Option<Rc<RefCell<TreeNode>>>) -> Option<Rc<RefCell<TreeNode>>> {
        if let Some(root_rc) = root.clone() {
            let tmp = Solution::prune_tree(root_rc.borrow().left.clone());
            root_rc.borrow_mut().left = tmp;
            let tmp = Solution::prune_tree(root_rc.borrow().right.clone());
            root_rc.borrow_mut().right = tmp;
            return if root_rc.borrow().val == 1 || root_rc.borrow().left.is_some() ||
                root_rc.borrow().right.is_some() { root } else { None };
        }
        None
    }
}


fn main() {
    Solution::prune_tree(None);
}
